import os
import cv2
import onnxruntime as ort
from PIL import Image, ImageDraw, ImageFont
import numpy as np
import sys

class LoadDetectOnnx:
    def __init__(self, model_path, confidence=0.7, nms_thresh=0.6):
        self.confidence = confidence
        self.nms_thresh = nms_thresh
        self.classes = {0: '火箭兵', 1: '工兵', 2: '普通士兵', 3: '指挥官'}
        self.color_palette = {
            0: (255, 0, 0),  # 火箭兵用红色
            1: (0, 255, 0),  # 工兵用绿色
            2: (0, 0, 255),  # 普通士兵用蓝色
            3: (255, 255, 0)  # 指挥官用黄色
        }
        self.providers = ['CPUExecutionProvider']
        self.session, self.model_inputs, self.input_width, self.input_height = self._init_model(model_path)
        
        # 加载中文字体
        font_path = "simhei.ttf"  # 确保该路径下有 simhei.ttf 字体文件
        self.font = ImageFont.truetype(font_path, 20) if os.path.exists(font_path) else None

    def _init_model(self, model_path):
        session = ort.InferenceSession(model_path, providers=self.providers)
        input_name = session.get_inputs()[0].name
        input_shape = session.get_inputs()[0].shape
        input_height, input_width = input_shape[2], input_shape[3]
        return session, input_name, input_width, input_height

    def _preprocess(self, image):
        resized = cv2.resize(image, (self.input_width, self.input_height))
        blob = cv2.dnn.blobFromImage(resized, scalefactor=1/255.0)
        return blob

    def run_inference(self, preprocessed_img):
        outputs = self.session.run(None, {self.model_inputs: preprocessed_img})
        return outputs

    def _draw(self, image, box, score, class_id):
        x1, y1, x2, y2 = box
        class_name = self.classes[class_id]
        color = self.color_palette[class_id]

        # 转换图像以使用 PIL 进行绘制
        pil_image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
        draw = ImageDraw.Draw(pil_image)

        # 绘制边界框
        draw.rectangle([x1, y1, x2, y2], outline=color, width=2)

        # 绘制类别和置信度文本
        text = f"{class_name} {score:.2f}"
        text_size = draw.textsize(text, font=self.font)
        text_background = (x1, y1 - text_size[1], x1 + text_size[0], y1)
        draw.rectangle(text_background, fill=color)
        draw.text((x1, y1 - text_size[1]), text, font=self.font, fill=(255, 255, 255))

        # 转换回 OpenCV 图像格式
        image = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
        return image

    def detect_object(self, img_src):
        image = cv2.imread(img_src)
        preprocessed_img = self._preprocess(image)
        outputs = self.run_inference(preprocessed_img)

        # 处理推理输出，提取边界框、置信度和类别
        boxes, scores, class_ids = [], [], []
        output_data = outputs[0]

        for det in output_data:
            score = det[4]
            class_id = int(det[5])
            if score >= self.confidence:
                x1, y1, x2, y2 = int(det[0]), int(det[1]), int(det[2]), int(det[3])
                boxes.append([x1, y1, x2, y2])
                scores.append(score)
                class_ids.append(class_id)

        # 非极大值抑制
        indices = cv2.dnn.NMSBoxes(boxes, scores, self.confidence, self.nms_thresh)
        filtered_indices = [i[0] for i in indices]

        # 类别计数
        class_counts = {class_name: 0 for class_name in self.classes.values()}
        for idx in filtered_indices:
            box = boxes[idx]
            score = scores[idx]
            class_id = class_ids[idx]
            class_name = self.classes[class_id]
            class_counts[class_name] += 1

            image = self._draw(image, box, score, class_id)

        return image, class_counts

def main(area):
    model_path = os.path.join(os.path.dirname(__file__), "best.onnx")
    infer = LoadDetectOnnx(model_path)
    img_list = os.listdir(os.path.join(os.path.dirname(__file__), 'image/photo'))

    # 记录总的类别数量
    total_class_counts = {class_name: 0 for class_name in infer.classes.values()}

    for img in img_list:
        if area in img:
            image_input_path = os.path.join(os.path.dirname(__file__), 'image/photo', img)
            output_image, class_counts = infer.detect_object(image_input_path)

            if output_image is not None:
                image_output_path = os.path.join(os.path.dirname(__file__), 'image/result', img)
                cv2.imwrite(image_output_path, output_image)

                for class_name, count in class_counts.items():
                    total_class_counts[class_name] += count

    # 输出区域检测结果
    print(f"区域 {area}:",end='')
    for class_name, count in total_class_counts.items():
        print(f",{class_name}: {count}人",end='')

if __name__ == "__main__":
    if len(sys.argv) > 1:
        main(sys.argv[1])
